#!/ usr / bin / env python3
#- * - coding : utf - 8 - * -

import numpy as np
import math
import sys
from get_ME import *
from sympy.physics.wigner import wigner_6j
from sympy.physics.wigner import wigner_9j

class Tensor_Decompose :
# Beyond Hartree-Fock in Nuclear Physics; a tale of correlations
#  M. Hjorth-Jensen
    def __init__(self, me):
        self.sp_state = me.sp_state
        self.ME = me
        self.me_K_dic = {}

    def phase(self,N):
        if(N%2 == 0):
            return 1
        else:
            return -1

    def ablsj(self,a,b,L,S,J):
        j2a = self.sp_state[a-1].j2
        j2b = self.sp_state[b-1].j2

        l_a = self.sp_state[a-1].l
        l_b = self.sp_state[b-1].l
        j_a = j2a/2.0
        j_b = j2b/2.0

        pf = (j2a+1)*(j2b+1)*(2*L+1)*(2*S+1)
        pf = math.sqrt(pf)

        ninJ = wigner_9j(l_a,0.5,j_a,l_b,0.5,j_b,L,S,J,prec=16)
        val = pf*ninJ

        return val


    def Tri(self, j1, j2, J12):
        val = 1
        if(J12 < abs(j1-j2)):
            val = -1
        if(J12 > abs(j1+j2)):
            val = -1
        return val

    def ab_group_nl(self,a,b):
        val = 1
        if(self.sp_state[a-1].n != self.sp_state[b-1].n):
            val = -1
        if(self.sp_state[a-1].l != self.sp_state[b-1].l):
            val = -1
        return val

    def build(self, K):
        for conf,val in self.ME.me_dic.items():
            a = int(conf[0])
            b = int(conf[2])
            c = int(conf[4])
            d = int(conf[6])
            J = int(conf[8])
            T = int(conf[10])

            val_k_me = self.cal(a,b,c,d,J,T,K)

            conf_k = conf + "-" + str(K)
            self.me_K_dic[conf_k] = val_k_me


    def print_me(self):
        for conf,val in self.me_K_dic.items():
            a = int(conf[0])
            b = int(conf[2])
            c = int(conf[4])
            d = int(conf[6])
            J = int(conf[8])
            T = int(conf[10])
            K = int(conf[12])

            print(a,b,c,d,"\t ",J,"\t ",T,"\t ",K,"\t " ,val)


    def cal_print(self,a,b,c,d,J,T,K):

        val = 0.0
        size_sp = self.ME.size_sp

        phase_pf = self.phase(J)
        pf = (2*K + 1)

        l_a = self.sp_state[a-1].l
        l_b = self.sp_state[b-1].l
        l_c = self.sp_state[c-1].l
        l_d = self.sp_state[d-1].l

        L_ab_min = abs(l_a - l_b)
        L_ab_max = abs(l_a + l_b)

        L_cd_min = abs(l_c - l_d)
        L_cd_max = abs(l_c + l_d)

        S_ab_min = 0
        S_ab_max = 1
        S_cd_min = 0
        S_cd_max = 1

        part_1 = 0.0
        for S_ab in range(S_ab_min,S_ab_max+1):
            for S_cd in range(S_cd_min,S_cd_max+1):
                if(self.Tri(S_ab,S_cd,K)<0):
                    continue
                for L_ab in range(L_ab_min,L_ab_max+1):
                    if(self.Tri(L_ab,S_ab,J)<0):
                        continue
                    for L_cd in range(L_cd_min,L_cd_max+1):
                        if(self.Tri(L_cd,S_cd,J)<0):
                            continue
                        if(self.Tri(L_ab,L_cd,K)<0):
                            continue
                        part_1_t =  wigner_6j(L_ab, S_ab, J, S_cd, L_cd, K,prec=16)
                        if(part_1_t == 0.0):
                            continue
                        part_1_t = part_1_t * self.ablsj(a,b,L_ab,S_ab,J)*self.ablsj(c,d,L_cd,S_cd,J)
                        if(part_1_t == 0.0):
                            continue

                        J_ab_min = abs( L_ab - S_ab )
                        J_ab_max = abs( L_ab + S_ab )
                        J_cd_min = abs( L_cd - S_cd )
                        J_cd_max = abs( L_cd + S_cd )

                        J_p_min = max(J_ab_min, J_cd_min)
                        J_p_max = min(J_ab_max, J_cd_max)

                        part_2 = 0.0
                        for J_p in range(J_p_min,J_p_max+1):
                            pf_2 = self.phase(J_p)*(2*J_p+1)
                            part_2_t = pf_2 * wigner_6j(L_ab, S_ab, J_p, S_cd, L_cd, K,prec=16)
                            if(part_2_t == 0.0):
                                continue
                            part_3 = 0.0
                            for i in range(size_sp):
                                i = i+1
                                if (self.ab_group_nl(i,a)<0):
                                    continue
                                for j in range(size_sp):
                                    j = j+1
                                    if (self.ab_group_nl(j,b)<0):
                                        continue
                                    pf_ijLSJp = self.ablsj(i,j,L_ab, S_ab, J_p)
                                    if(pf_ijLSJp == 0.0):
                                        continue
                                    for k in range(size_sp):
                                        k = k+1
                                        if (self.ab_group_nl(k,c)<0):
                                            continue
                                        for l in range(size_sp):
                                            l = l+1
                                            if (self.ab_group_nl(l,d)<0):
                                                continue
                                            pf_klLSJp = self.ablsj(k,l,L_cd, S_cd, J_p)
                                            if(pf_klLSJp == 0.0):
                                                continue

                                            val_me = self.ME.get_ME(i,j,k,l,J_p,T)
                                            print("--- ",i,j,k,l,"\t ",J_p,"\t ",T)
                                            print("\t\t ",val_me, pf_ijLSJp, pf_klLSJp)
                                            part_3_t = val_me*pf_ijLSJp*pf_klLSJp

                                            part_3 = part_3 + part_3_t
                                        #l
                                    #k
                                #j
                            #i
                            part_2 = part_2 + part_2_t*part_3
                        #J_p
                        part_1 = part_1 + part_1_t * part_2
                    #L_cd
                #L_ab
            #S_cd
        #S_ab
        val = phase_pf * pf * part_1
        return val


    def cal(self,a,b,c,d,J,T,K):

        val = 0.0
        size_sp = self.ME.size_sp

        phase_pf = self.phase(J)
        pf = (2*K + 1)

        l_a = self.sp_state[a-1].l
        l_b = self.sp_state[b-1].l
        l_c = self.sp_state[c-1].l
        l_d = self.sp_state[d-1].l

        L_ab_min = abs(l_a - l_b)
        L_ab_max = abs(l_a + l_b)

        L_cd_min = abs(l_c - l_d)
        L_cd_max = abs(l_c + l_d)

        S_ab_min = 0
        S_ab_max = 1
        S_cd_min = 0
        S_cd_max = 1

        part_1 = 0.0
        for S_ab in range(S_ab_min,S_ab_max+1):
            for S_cd in range(S_cd_min,S_cd_max+1):
                if(self.Tri(S_ab,S_cd,K)<0):
                    continue
                for L_ab in range(L_ab_min,L_ab_max+1):
                    if(self.Tri(L_ab,S_ab,J)<0):
                        continue
                    for L_cd in range(L_cd_min,L_cd_max+1):
                        if(self.Tri(L_cd,S_cd,J)<0):
                            continue
                        if(self.Tri(L_ab,L_cd,K)<0):
                            continue
                        part_1_t =  wigner_6j(L_ab, S_ab, J, S_cd, L_cd, K,prec=16)
                        if(part_1_t == 0.0):
                            continue
                        part_1_t = part_1_t * self.ablsj(a,b,L_ab,S_ab,J)*self.ablsj(c,d,L_cd,S_cd,J)
                        if(part_1_t == 0.0):
                            continue

                        J_ab_min = abs( L_ab - S_ab )
                        J_ab_max = abs( L_ab + S_ab )
                        J_cd_min = abs( L_cd - S_cd )
                        J_cd_max = abs( L_cd + S_cd )

                        J_p_min = max(J_ab_min, J_cd_min)
                        J_p_max = min(J_ab_max, J_cd_max)

                        part_2 = 0.0
                        for J_p in range(J_p_min,J_p_max+1):
                            pf_2 = self.phase(J_p)*(2*J_p+1)
                            part_2_t = pf_2 * wigner_6j(L_ab, S_ab, J_p, S_cd, L_cd, K,prec=16)
                            if(part_2_t == 0.0):
                                continue
                            part_3 = 0.0
                            for i in range(size_sp):
                                i = i+1
                                if (self.ab_group_nl(i,a)<0):
                                    continue
                                for j in range(size_sp):
                                    j = j+1
                                    if (self.ab_group_nl(j,b)<0):
                                        continue
                                    pf_ijLSJp = self.ablsj(i,j,L_ab, S_ab, J_p)
                                    if(pf_ijLSJp == 0.0):
                                        continue
                                    for k in range(size_sp):
                                        k = k+1
                                        if (self.ab_group_nl(k,c)<0):
                                            continue
                                        for l in range(size_sp):
                                            l = l+1
                                            if (self.ab_group_nl(l,d)<0):
                                                continue
                                            pf_klLSJp = self.ablsj(k,l,L_cd, S_cd, J_p)
                                            if(pf_klLSJp == 0.0):
                                                continue
                                            val_me = self.ME.get_ME(i,j,k,l,J_p,T)
                                            part_3_t = val_me*pf_ijLSJp*pf_klLSJp

                                            part_3 = part_3 + part_3_t
                                        #l
                                    #k
                                #j
                            #i
                            part_2 = part_2 + part_2_t*part_3
                        #J_p
                        part_1 = part_1 + part_1_t * part_2
                    #L_cd
                #L_ab
            #S_cd
        #S_ab
        val = phase_pf * pf * part_1
        return val
